
from copy import copy, deepcopy
import torch
import torch.nn as nn

from fast3r.dust3r.losses import MultiLoss, Criterion

class BaseCriterion(nn.Module):
    def __init__(self, reduction="mean"):
        super().__init__()
        self.reduction = reduction

class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images"""

    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool = nn.AvgPool2d(3, 1)
        self.mu_y_pool = nn.AvgPool2d(3, 1)
        self.sig_x_pool = nn.AvgPool2d(3, 1)
        self.sig_y_pool = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01**2
        self.C2 = 0.03**2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x = self.sig_x_pool(x**2) - mu_x**2
        sigma_y = self.sig_y_pool(y**2) - mu_y**2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x**2 + mu_y**2 + self.C1) * (sigma_x + sigma_y + self.C2)

        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)
class RGBLoss(Criterion, MultiLoss):
    def __init__(self, criterion, ssim_weight: float = 0.0):
        super().__init__(criterion)
        if ssim_weight > 0:
            self.ssim = SSIM()
        else:
            self.ssim = None

    def img_loss(self, a, b):
        return self.criterion(a, b)

    def compute_loss(self, gts, preds, **kw):
        gt_rgbs = [gt["img"].permute(0, 2, 3, 1) for gt in gts]
        pred_rgbs = [pred["rgb"] for pred in preds]
        ls = [
            self.img_loss(pred_rgb, gt_rgb)
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs)
        ]
        rgb_loss = sum(ls) / len(ls)

        details = {}
        # self_name = type(self).__name__
        for i, l in enumerate(ls):
            details[f"RGBLoss_rgb/{i+1}"] = float(l)
            # details[f"pred_rgb_{i+1}"] = pred_rgbs[i]
        details['RGBLoss_avg_rgb']  = float(rgb_loss)

        if self.ssim:
            # perceptual SSIM losses
            ssim_losses = []
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs):
                # permute to (B,C,H,W)
                p = pred_rgb.permute(0,3,1,2)
                g = gt_rgb.permute(0,3,1,2)
                ssim_losses.append(self.ssim(p, g))
            avg_ssim_loss = sum(ssim_losses) / len(ssim_losses)

            for i, l in enumerate(ssim_losses, 1):
                details[f"RGBLoss_ssim/{i}"] = float(l)
            details['RGBLoss_avg_ssim'] = float(avg_ssim_loss)

            rgb_loss = rgb_loss + self.ssim_weight * avg_ssim_loss

        return rgb_loss, details

